import bnlearn as bn
import data_utils
import pandas as pd
from colorama import Fore


def eval_test():
    # data = data_utils.read_data_file(
    # data_utils.DataName.Task1, data_utils.DataName.Heterozygous
    # )
    data = data_utils.read_Dream4_data()
    #     model1 = BayesianNetwork([("G1", "G2"), ("G1", "G3")])
    # model2 = BayesianNetwork([("G1", "G5"), ("G2", "G3")])

    edges_1 = [("G1", "G2"), ("G1", "G3")]
    scores_1 = get_scores(edges_1, data)
    print(f"model1 score = {Fore.GREEN}{scores_1}{Fore.RESET}")
    edges_2 = [("G1", "G5"), ("G2", "G3")]
    scores_2 = get_scores(edges_2, data)
    print(f"model2 score = {Fore.GREEN}{scores_2}{Fore.RESET}")

    edges = [
        ("G2", "G1"),
        ("G2", "G3"),
        ("G3", "G4"),
        ("G9", "G4"),
        ("G3", "G5"),
        ("G8", "G5"),
        ("G9", "G5"),
        ("G3", "G6"),
        ("G3", "G7"),
        ("G8", "G7"),
        ("G10", "G7"),
    ]
    score = get_scores(edges, data)
    print(f"gold model score = {Fore.GREEN}{score}{Fore.RESET}")

    back_edges = [
        ("G1", "G2"),
        ("G3", "G2"),
        ("G4", "G3"),
        ("G4", "G9"),
        ("G5", "G3"),
        ("G5", "G8"),
        ("G5", "G9"),
        ("G6", "G3"),
        ("G7", "G3"),
        ("G7", "G8"),
        ("G7", "G10"),
    ]
    score = get_scores(back_edges, data)
    print(f"back model score = {Fore.GREEN}{score}{Fore.RESET}")


def get_scores(edges, data):
    data = pd.DataFrame(data, columns=data_utils.DataName.ColumnNames)
    DGA = bn.make_DAG(edges, methodtype="bayes")
    model = DGA
    model["config"] = {"method": "bayes"}
    score_get = bn.structure_scores(model, data, scoring_method="bic")
    return score_get


if __name__ == "__main__":
    eval_test()
